#!/usr/bin/env python3
"""
改进的解密脚本 - 基于已发现的模式
"""

from collections import Counter, defaultdict
import re

def load_data():
    """加载数据"""
    with open('decoded_result.bin', 'rb') as f:
        decoded = f.read()
    
    with open('partial_decrypt.bin', 'rb') as f:
        partial = f.read()
    
    return decoded, partial

def analyze_known_patterns():
    """分析已知的模式"""
    print("🔍 === 分析已知模式 ===\n")
    
    decoded, partial = load_data()
    
    # 从部分解密中我们知道 "model" 正确解密了
    # 让我们找出对应的加密字节
    model_positions = []
    model_pattern = b"model"
    
    # 在部分解密中找到 "model" 的位置
    pos = 0
    while True:
        pos = partial.find(model_pattern, pos)
        if pos == -1:
            break
        model_positions.append(pos)
        pos += 1
    
    print(f"找到 'model' {len(model_positions)} 次，位置: {model_positions[:10]}")
    
    # 获取 "model" 对应的加密字节
    if model_positions:
        # 检查第一个 model 的加密字节
        pos = model_positions[0]
        encrypted_model = decoded[pos:pos+5]
        print(f"\n'model' 的加密字节: {encrypted_model.hex()}")
        
        # 验证映射
        mapping = {}
        for i, char in enumerate(b"model"):
            mapping[encrypted_model[i]] = char
        
        print("\n确认的字符映射:")
        for enc, dec in mapping.items():
            print(f"  0x{enc:02x} -> '{chr(dec)}'")
    
    return mapping

def expand_mapping(decoded, initial_mapping):
    """扩展字符映射表"""
    print("\n\n🔧 === 扩展映射表 ===")
    
    mapping = initial_mapping.copy()
    
    # 已经确认的映射
    # 0x69 -> " (引号)
    # 0x01 -> : (冒号)
    # 0x90 -> 空格
    # 从 "model" 我们得到更多映射
    
    # 常见的JSON字段名
    common_fields = [
        b"questionText",
        b"conversationId", 
        b"sessionId",
        b"userId",
        b"content",
        b"message",
        b"response",
        b"timestamp",
        b"status",
        b"type",
        b"data",
        b"result",
        b"error",
        b"code",
        b"name",
        b"value",
        b"id"
    ]
    
    # 基于已知的部分映射，尝试找到更多
    quote = 0x69  # 引号
    
    print("\n查找可能的字段名:")
    
    # 查找 "xxx": 模式
    field_patterns = []
    i = 0
    while i < len(decoded) - 3:
        if decoded[i] == quote:
            # 找下一个引号
            j = i + 1
            while j < len(decoded) and decoded[j] != quote:
                j += 1
            
            if j < len(decoded) - 1 and decoded[j+1] == 0x01:  # 冒号
                # 找到一个字段
                field_bytes = decoded[i+1:j]
                if 3 <= len(field_bytes) <= 20:  # 合理的字段名长度
                    field_patterns.append(field_bytes)
                i = j + 1
            else:
                i += 1
        else:
            i += 1
    
    # 统计字段模式
    field_counter = Counter(field_patterns)
    common_encrypted_fields = field_counter.most_common(20)
    
    print(f"\n最常见的加密字段 (共{len(field_counter)}个):")
    for field, count in common_encrypted_fields[:10]:
        # 尝试部分解密
        decrypted = bytearray()
        unknown_count = 0
        for b in field:
            if b in mapping:
                decrypted.append(mapping[b])
            else:
                decrypted.append(ord('?'))
                unknown_count += 1
        
        print(f"  {field.hex()}: {count}次 -> '{bytes(decrypted).decode('ascii', errors='ignore')}' ({unknown_count}个未知)")
        
        # 尝试匹配已知字段名
        if unknown_count <= 3:  # 只有少数未知字符
            for known_field in common_fields:
                if len(known_field) == len(field):
                    # 检查已知部分是否匹配
                    match = True
                    for k in range(len(field)):
                        if field[k] in mapping and mapping[field[k]] != known_field[k]:
                            match = False
                            break
                    
                    if match:
                        print(f"    可能是: '{known_field.decode()}'")
                        # 添加新映射
                        for k in range(len(field)):
                            if field[k] not in mapping:
                                mapping[field[k]] = known_field[k]
                                print(f"      新映射: 0x{field[k]:02x} -> '{chr(known_field[k])}'")
    
    return mapping

def apply_full_mapping(decoded, mapping):
    """应用完整的映射表"""
    print("\n\n📝 === 应用映射表 ===")
    
    print(f"当前映射表有 {len(mapping)} 个字符")
    
    # 应用映射
    decrypted = bytearray()
    unmapped_bytes = set()
    
    for byte in decoded:
        if byte in mapping:
            decrypted.append(mapping[byte])
        else:
            decrypted.append(byte)
            unmapped_bytes.add(byte)
    
    # 统计
    mapped_count = len(decoded) - sum(1 for b in decoded if b not in mapping)
    print(f"已映射: {mapped_count}/{len(decoded)} ({mapped_count/len(decoded)*100:.1f}%)")
    print(f"未映射的唯一字节: {len(unmapped_bytes)}个")
    
    # 保存结果
    with open('improved_decrypt.txt', 'wb') as f:
        f.write(decrypted)
    
    print("\n解密结果已保存到 improved_decrypt.txt")
    
    # 显示预览
    preview = decrypted[:500]
    # 将不可打印字符替换为 ?
    printable_preview = bytearray()
    for b in preview:
        if 32 <= b <= 126:
            printable_preview.append(b)
        else:
            printable_preview.append(ord('?'))
    
    print(f"\n预览 (前500字符):")
    print(bytes(printable_preview).decode('ascii'))
    
    # 尝试提取JSON结构
    extract_json_structure(decrypted)
    
    return decrypted, unmapped_bytes

def extract_json_structure(decrypted):
    """尝试提取JSON结构"""
    print("\n\n🎯 === 提取JSON结构 ===")
    
    # 查找JSON特征
    json_start = decrypted.find(b'{')
    json_end = decrypted.rfind(b'}')
    
    if json_start >= 0 and json_end > json_start:
        json_content = decrypted[json_start:json_end+1]
        
        # 清理不可打印字符
        cleaned = bytearray()
        for b in json_content:
            if 32 <= b <= 126:
                cleaned.append(b)
            else:
                cleaned.append(ord(' '))  # 用空格替代
        
        # 尝试格式化
        text = bytes(cleaned).decode('ascii', errors='ignore')
        
        # 提取可识别的字段
        print("\n可识别的JSON字段:")
        
        # 查找 "fieldname": 模式
        field_pattern = r'"([a-zA-Z_]+)":'
        fields = re.findall(field_pattern, text)
        
        if fields:
            field_counter = Counter(fields)
            for field, count in field_counter.most_common():
                print(f"  {field}: {count}次")
        
        # 查找值
        print("\n部分内容示例:")
        # 查找 "model": "xxx" 模式
        model_pattern = r'"model":\s*"([^"]*)"'
        models = re.findall(model_pattern, text)
        if models:
            print(f"  model 值: {set(models)}")
        
        # 保存清理后的JSON
        with open('cleaned_json.txt', 'w') as f:
            f.write(text)
        print("\n清理后的JSON已保存到 cleaned_json.txt")

def suggest_remaining_mappings(decoded, unmapped_bytes, current_mapping):
    """建议剩余的映射"""
    print("\n\n💡 === 映射建议 ===")
    
    # 分析未映射字节的上下文
    context_analysis = defaultdict(lambda: defaultdict(int))
    
    for i in range(1, len(decoded)-1):
        if decoded[i] in unmapped_bytes:
            # 获取前后字符
            prev_byte = decoded[i-1]
            next_byte = decoded[i+1]
            
            prev_char = chr(current_mapping.get(prev_byte, ord('?')))
            next_char = chr(current_mapping.get(next_byte, ord('?')))
            
            context = f"{prev_char}_{next_char}"
            context_analysis[decoded[i]][context] += 1
    
    print("未映射字节的常见上下文:")
    
    # 基于上下文的猜测
    suggestions = {}
    
    for byte_val, contexts in sorted(context_analysis.items(), 
                                   key=lambda x: sum(x[1].values()), 
                                   reverse=True)[:10]:
        total = sum(contexts.values())
        print(f"\n0x{byte_val:02x} (出现{total}次):")
        
        for context, count in sorted(contexts.items(), 
                                    key=lambda x: x[1], 
                                    reverse=True)[:3]:
            print(f"  {context}: {count}次")
            
            # 基于上下文猜测
            if '"_"' in context:
                suggestions[byte_val] = "字段名中的字符 (a-z)"
            elif ':_' in context and '_"' in context:
                suggestions[byte_val] = "值 (可能是数字或字母)"
            elif ':_' in context and '_ ' in context:
                suggestions[byte_val] = "值的开始 (可能是 [ 或 {)"
    
    if suggestions:
        print("\n基于上下文的猜测:")
        for byte_val, guess in suggestions.items():
            print(f"  0x{byte_val:02x}: {guess}")

if __name__ == "__main__":
    # 分析已知模式
    initial_mapping = analyze_known_patterns()
    
    # 加载数据
    decoded, _ = load_data()
    
    # 扩展映射
    expanded_mapping = expand_mapping(decoded, initial_mapping)
    
    # 应用映射
    decrypted, unmapped = apply_full_mapping(decoded, expanded_mapping)
    
    # 建议剩余映射
    suggest_remaining_mappings(decoded, unmapped, expanded_mapping)
    
    print("\n\n✅ === 总结 ===")
    print(f"成功映射了 {len(expanded_mapping)} 个字符")
    print("查看以下文件了解进展:")
    print("- improved_decrypt.txt: 改进的解密结果")
    print("- cleaned_json.txt: 清理后的JSON内容")
    print("\n下一步:")
    print("1. 根据上下文继续完善字符映射")
    print("2. 或使用更多的已知明文")
    print("3. 或通过动态调试获取完整映射表")